{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此模型采用bert[BERT for Joint Intent Classification and Slot Filling](https://arxiv.org/pdf/1902.10909.pdf)进行意图识别与槽填充。\n",
    "结构如下："
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![model](img/bert.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "从上可知：\n",
    "\n",
    "    1.意图识别采用[cls]的输出进行识别\n",
    "    \n",
    "    2.槽填充直接输出对应的结果进行序列标注，这里不使用mlm中的mask\n",
    "    \n",
    "步骤：\n",
    "\n",
    "    1.输入是单句，[cls]用于最后的意图识别, [sep]作为句子的最后部分\n",
    "    \n",
    "    2.再加入位置信息\n",
    "    \n",
    "    3.经过embedding -> transformer-encoder -> 输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from torchtext import data, datasets\n",
    "import pandas as pd\n",
    "import pickle\n",
    "\n",
    "base_dir = os.path.abspath(os.path.join(os.getcwd(), \"..\"))\n",
    "atis_data = os.path.join(base_dir, 'atis')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "build train and val dataset\n",
    "'''\n",
    "    \n",
    "tokenize = lambda s:s.split()\n",
    "\n",
    "SOURCE = data.Field(sequential=True, tokenize=tokenize,\n",
    "                    lower=True, use_vocab=True,\n",
    "                    init_token='<cls>', eos_token='<sep>',\n",
    "                    pad_token='<pad>', unk_token='<unk>',\n",
    "                    batch_first=True, fix_length=50,\n",
    "                    include_lengths=True) #include_lengths=True为方便之后使用torch的pack_padded_sequence\n",
    "\n",
    "TARGET = data.Field(sequential=True, tokenize=tokenize,\n",
    "                    lower=True, use_vocab=True,\n",
    "                    init_token='<cls>', eos_token='<sep>',\n",
    "                    pad_token='<pad>', unk_token='<unk>',\n",
    "                    batch_first=True, fix_length=50,\n",
    "                    include_lengths=True) #include_lengths=True为方便之后使用torch的pack_padded_sequence\n",
    "LABEL = data.Field(\n",
    "                sequential=False,\n",
    "                use_vocab=True)\n",
    "\n",
    "train, val = data.TabularDataset.splits(\n",
    "                                        path=atis_data,\n",
    "                                        skip_header=True,\n",
    "                                        train='atis.train.csv',\n",
    "                                        validation='atis.test.csv',\n",
    "                                        format='csv',\n",
    "                                        fields=[('index', None), ('intent', LABEL), ('source', SOURCE), ('target', TARGET)])\n",
    "\n",
    "SOURCE.build_vocab(train, val)\n",
    "TARGET.build_vocab(train, val)\n",
    "LABEL.build_vocab(train, val)\n",
    "\n",
    "train_iter, val_iter = data.Iterator.splits(\n",
    "                                            (train, val),\n",
    "                                            batch_sizes=(64, len(val)), # 训练集设置为32,验证集整个集合用于测试\n",
    "                                            shuffle=True,\n",
    "                                            sort_within_batch=True, #为true则一个batch内的数据会按sort_key规则降序排序\n",
    "                                            sort_key=lambda x: len(x.source)) #这里按src的长度降序排序，主要是为后面pack,pad操作)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save source words\n",
    "source_words_path = os.path.join(os.getcwd(), 'source_words.pkl')\n",
    "with open(source_words_path, 'wb') as f_source_words:\n",
    "    pickle.dump(SOURCE.vocab, f_source_words)\n",
    "\n",
    "# save target words\n",
    "target_words_path = os.path.join(os.getcwd(), 'target_words.pkl')\n",
    "with open(target_words_path, 'wb') as f_target_words:\n",
    "    pickle.dump(TARGET.vocab, f_target_words)\n",
    "    \n",
    "# save label words\n",
    "label_words_path = os.path.join(os.getcwd(), 'label_words.pkl')\n",
    "with open(label_words_path, 'wb') as f_label_words:\n",
    "    pickle.dump(LABEL.vocab, f_label_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<cls>, <sep>\n"
     ]
    }
   ],
   "source": [
    "print('{}, {}'.format(SOURCE.init_token, SOURCE.eos_token))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import math\n",
    "from apex import amp\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "SEED = 1234\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic=True\n",
    "torch.backends.cudnn.benchmark=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, max_length=100):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.tok_embedding = nn.Embedding(input_dim, hid_dim)\n",
    "        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n",
    "        \n",
    "        # 多层encoder\n",
    "        self.layers = nn.ModuleList([EncoderLayer(hid_dim, n_heads, pf_dim, dropout) for _ in range(n_layers)])\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(DEVICE)\n",
    "\n",
    "    def forward(self, src, src_mask):\n",
    "        #src:[batch_size, src_len]\n",
    "        #src_mask:[batch_size, 1, 1, src_len]\n",
    "        \n",
    "        batch_size = src.shape[0]\n",
    "        src_len = src.shape[1]\n",
    "        #位置信息\n",
    "        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(DEVICE)\n",
    "        #token编码+位置编码\n",
    "        src = self.dropout((self.tok_embedding(src) * self.scale) + self.pos_embedding(pos)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            src = layer(src, src_mask) #[batch_size, src_len, hid_dim]\n",
    "        \n",
    "        return src\n",
    "            \n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, n_heads, pf_dim, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout)\n",
    "        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, pf_dim, dropout)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src, src_mask):\n",
    "        # src:[batch_size, src_len, hid_dim]\n",
    "        # src_mask:[batch_size, 1, 1, src_len]\n",
    "        \n",
    "        # 1.经过多头attetnion后，再经过add+norm\n",
    "        # self-attention\n",
    "        _src = self.self_attention(src, src, src, src_mask)\n",
    "        \n",
    "        src = self.self_attn_layer_norm(src + self.dropout(_src)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 2.经过一个前馈网络后，再经过add+norm\n",
    "        _src = self.positionwise_feedforward(src)\n",
    "        \n",
    "        src = self.ff_layer_norm(src + self.dropout(_src)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        return src     \n",
    "\n",
    "class MultiHeadAttentionLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, n_heads, dropout):\n",
    "        super(MultiHeadAttentionLayer, self).__init__()\n",
    "        assert hid_dim % n_heads == 0\n",
    "        \n",
    "        self.hid_dim = hid_dim\n",
    "        self.n_heads = n_heads\n",
    "        self.head_dim = hid_dim // n_heads\n",
    "        \n",
    "        self.fc_q = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_k = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_v = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_o = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(DEVICE)\n",
    "        \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        batch_size = query.shape[0]\n",
    "        \n",
    "        # query:[batch_size, query_len, hid_dim]\n",
    "        # key:[batch_size, query_len, hid_dim]\n",
    "        # value:[batch_size, query_len, hid_dim]\n",
    "        \n",
    "        Q = self.fc_q(query)\n",
    "        K = self.fc_k(key)\n",
    "        V = self.fc_v(value)\n",
    "        \n",
    "        Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3) # [batch_size, query_len, n_heads, head_dim]\n",
    "        K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        \n",
    "        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale # [batch_size, n_heads, query_len, key_len]\n",
    "        \n",
    "        if mask is not None:\n",
    "            energy = energy.mask_fill(mask == 0, -1e10)\n",
    "        \n",
    "        attention = torch.softmax(energy, dim=-1) # [batch_size, n_heads, query_len, key_len]\n",
    "        \n",
    "        x = torch.matmul(self.dropout(attention), V) # [batch_size, n_heads, query_len, head_dim]\n",
    "        \n",
    "        x = x.permute(0, 2, 1, 3).contiguous() # [batch_size, query_len, n_heads, head_dim]\n",
    "        \n",
    "        x = x.view(batch_size, -1, self.hid_dim) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        x = self.fc_o(x) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        return x\n",
    "        \n",
    "class PositionwiseFeedforwardLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, pf_dim, dropout):\n",
    "        super(PositionwiseFeedforwardLayer, self).__init__()\n",
    "        self.fc_1 = nn.Linear(hid_dim, pf_dim)\n",
    "        self.fc_2 = nn.Linear(pf_dim, hid_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.gelu = nn.GELU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x:[batch_size, seq_len, hid_dim]\n",
    "        \n",
    "        x = self.dropout(self.gelu(self.fc_1(x))) # [batch_size, seq_len, pf_dim]\n",
    "        x = self.fc_2(x) # [batch_size, seq_len, hid_dim]\n",
    "        \n",
    "        return x\n",
    "\n",
    "class BERT(nn.Module):\n",
    "    def __init__(self, input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, slot_size, intent_size, src_pad_idx):\n",
    "        super(BERT, self).__init__()\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.encoder = Encoder(input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout)\n",
    "        self.gelu = nn.GELU()\n",
    "        \n",
    "        self.fc = nn.Sequential(nn.Linear(hid_dim, hid_dim), nn.Dropout(dropout), nn.Tanh())\n",
    "        self.intent_out = nn.Linear(hid_dim, intent_size)\n",
    "        self.linear = nn.Linear(hid_dim, hid_dim)\n",
    "       \n",
    "        embed_weight = self.encoder.tok_embedding.weight\n",
    "        self.slot_out = nn.Linear(hid_dim, slot_size, bias=False)\n",
    "        # 此处slot_out共享tok_embedding的权重\n",
    "        self.slot_out.weight = embed_weight\n",
    "    \n",
    "    def make_src_mask(self, src):\n",
    "        # src: [batch_size, src_len]\n",
    "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2) # [batch_size, 1, 1, src_len]\n",
    "        \n",
    "    def forward(self, src):\n",
    "        src_mask = self.make_src_mask(src)\n",
    "        encoder_out = self.encoder(src, src_mask) #[batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 拿到[cls] token进行意图分类\n",
    "        cls_hidden = self.fc(encoder_out[:, 0]) # [batch_size, hid_dim]\n",
    "        intent_output = self.intent_out(cls_hidden) # [batch_size, intent_size]\n",
    "        \n",
    "        # 排除cls进行slot预测\n",
    "        other_hidden = self.gelu(self.linear(encoder_out[:,1:])) # [batch_sze, src_len-1, hid_dim]\n",
    "        slot_output = self.slot_out(other_hidden) # [batch_size, src_len-1, slot_size]\n",
    "        return intent_output, slot_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_layers = 6 # transformer-encoder层数\n",
    "n_heads = 12 # 多头self-attention\n",
    "hid_dim =768 \n",
    "dropout = 0.5\n",
    "pf_dim = 768 * 4 \n",
    "\n",
    "input_dim = len(SOURCE.vocab)\n",
    "slot_size = len(TARGET.vocab) # slot size\n",
    "intent_size = len(LABEL.vocab) # intent size\n",
    "src_pad_idx = SOURCE.vocab.stoi[SOURCE.pad_token]\n",
    "\n",
    "model = BERT(input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, slot_size, intent_size, src_pad_idx).to(DEVICE)\n",
    "\n",
    "# 优化函数\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "# 损失函数(slot)\n",
    "loss_slot = nn.CrossEntropyLoss(ignore_index=src_pad_idx)\n",
    "\n",
    "# 定义损失函数(意图识别)\n",
    "loss_intent = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练\n",
    "def train(model, iterator, optimizer, loss_slot, loss_intent, clip):\n",
    "    \n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for i, batch in enumerate(iterator):\n",
    "        \n",
    "        src, _ = batch.source  # src=[batch_size, seq_len]，这里batch.src返回src和src的长度，因为在使用torchtext.Field时设置include_lengths=True\n",
    "        trg, _ = batch.target  # trg=[batch_size, seq_len]\n",
    "        label = batch.intent # [batch_size]\n",
    "        src = src.to(DEVICE)\n",
    "        trg = trg.to(DEVICE)\n",
    "        label = label.to(DEVICE)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        intent_output, slot_output = model(src) # [batch_size, intent_size]; [batch_size, trg_len-1, slot_size]\n",
    "        \n",
    "        # 1.计算slot loss\n",
    "        slot_output_dim = slot_output.shape[-1]\n",
    "        \n",
    "        slot_output = slot_output.reshape(-1, slot_output_dim) # [batch_size * (trg_len-1), slot_output_dim]\n",
    "        \n",
    "        trg = trg[:,1:].contiguous().view(-1) # [batch_size * (trg_len-1)]\n",
    "        \n",
    "        # 1.计算slot loss\n",
    "        loss1 = loss_slot(slot_output, trg)\n",
    "        \n",
    "        # 2.计算intent loss\n",
    "        loss2 = loss_intent(intent_output, label)\n",
    "        \n",
    "        # 3.联合slot loss + intent loss\n",
    "        loss = loss1 + loss2\n",
    "        \n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "    \n",
    "    return epoch_loss / len(iterator)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# val loss\n",
    "def evaluate(model, iterator, loss_slot, loss_intent):\n",
    "    model.eval()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for i, batch in enumerate(iterator):\n",
    "            src, _ = batch.source  # src=[batch_size, seq_len]\n",
    "            trg, _ = batch.target  # trg=[batch_size, seq_len]\n",
    "            label = batch.intent\n",
    "            src = src.to(DEVICE)\n",
    "            trg = trg.to(DEVICE)\n",
    "            label = label.to(DEVICE)\n",
    "            \n",
    "            intent_output, slot_output = model(src) # [batch_size, intent_dim]; [batch_size, trg_len-1, slot_size]\n",
    "            \n",
    "            # 1.计算slot loss\n",
    "            slot_output_dim = slot_output.shape[-1]\n",
    "\n",
    "            slot_output = slot_output.reshape(-1, slot_output_dim) # [batch_size * (trg_len-1), slot_output_dim]\n",
    "\n",
    "            trg = trg[:,1:].contiguous().view(-1) # [batch_size * (trg_len-1)]\n",
    "\n",
    "            loss1 = loss_slot(slot_output, trg)\n",
    "\n",
    "            # 2.计算intent loss\n",
    "            loss2 = loss_intent(intent_output, label)\n",
    "\n",
    "            loss = loss1 + loss2\n",
    "        \n",
    "            \n",
    "            epoch_loss += loss.item()\n",
    "            \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 0m 22s\n",
      "\tTrain Loss: 37.131 | Train PPL: 13362582122889738.000\n",
      "\t Val. Loss: 8.320 |  Val. PPL: 4104.540\n",
      "Epoch: 02 | Time: 0m 21s\n",
      "\tTrain Loss: 3.868 | Train PPL:  47.841\n",
      "\t Val. Loss: 5.999 |  Val. PPL: 402.855\n"
     ]
    }
   ],
   "source": [
    "n_epochs = 100 # 迭代次数\n",
    "clip = 0.1 # 梯度裁剪\n",
    "\n",
    "model_path = os.path.join(os.getcwd(), \"model.h5\")\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_iter, optimizer, loss_slot, loss_intent, clip)\n",
    "    valid_loss = evaluate(model, val_iter, loss_slot, loss_intent)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time) # 每个epoch花费的时间\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), model_path)\n",
    "        \n",
    "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}